#!/usr/bin/zsh
#############################################
# AutoScreen
# (C) 2010 Wes Alvaro
# Screens your sessions automatically, and
# prevents multiple nested screen sessions
# given a NFS shared folder.
# 
# Usage:
#  source autoScreen
#
#############################################

export sharedFolder=`echo $HOME` # The network folder
export ASext=`whoami`-AS

#############################################

PATH=$PATH:/sbin/ # For ipconfig

getMyIp() {
    export IP=`ifconfig  | grep 'inet addr:'| grep -v '127.0.0.1' | cut -d: -f2 | awk '{ print $1}'`
}
getDomainName() {
    export DOMAIN_NAME=`nslookup $1 | tail -n+4 | head -n1 | egrep -o "[0-9a-zA-Z.]+$" | egrep -o "^[a-zA-Z0-9]+"`
}
myScreen() {
    echo $STY;
}
myScreens() {
    screen -ls | grep "${ASext}"
}
autoScreenMe() {
  if [[ "$STY" == "" ]]; then # best I can tell, we're not in a screen, yet
    touch .screenHere
    local origin=`echo $SSH_CLIENT | cut -d' ' -f1`
    echo origin > .screenHere
    local screenOn="$sharedFolder/.screenOn-$origin"
    if [[ (! -e $screenOn)  ||  ($screenOn -ot .screenHere) ]]; then # we're not nested
      screen -ls > .screenResult
      local total=`cat .screenResult | grep "${ASext}" | wc -l` # all AS screens
      local s=$[`cat .screenResult | grep "${ASext}" | grep -iv "attached" | wc -l`] # only detached
      if [[ $total == 0 ]]; then # start a new screen
        echo "Starting new screen session as no unattached screen exists"
        screen -S ${ASext}
      elif [[$s == 0]] ; then # no unattached screens exist, clone? steal?
        echo "All attached! Clone/Steal an existing screen (1...$total)? (default: new screen)"
        read clone
        if [[ $clone == "" ]]; then # start a new screen
          echo "Starting new screen session as no unattached screen exists"
          screen -S ${ASext}
        else # clone a screen
          echo "Steal (s)? (default: clone)"
          read steal
          local c=$[ $total - ($clone - 1) ];
          local r=`screen -ls | grep "${ASext}" | tail -n $c | head -n 1 | awk -F. '{printf "%d\n", $1}'`
          if [[ $steal == "s" ]]; then # detach the remote screens
            echo "Stealing $clone of $total screen sessions"
            screen -D -r $r.${ASext}
          else
            echo "Cloning $clone of $total screen sessions"
            screen -x $r.${ASext}
          fi
        fi
      elif [[ $s == 1 ]]; then # there is only one detached screen. Take it!
        echo "Resuming the only unattached screen session ($total total)"
        local r=`screen -ls | grep "${ASext}" | grep -iv "attached" | head -n 1 | awk -F. '{printf "%d\n", $1}'`
        screen -r $r.${ASext}
      else # there are several unattached screens
        echo "Resuming the most recent of $s unattached screen sessions ($total total)"
        local r=`screen -ls | grep "${ASext}" | grep -iv "attached" | head -n 1 | awk -F. '{printf "%d\n", $1}'`
        screen -r $r.${ASext}
      fi
    else # You are already in a screen
      export STY=`cat $screenOn | cut -d: -f1`
      export WINDOW=`cat $screenOn | cut -d: -f2`
      echo "I will not automatically nest you in screens: ($WINDOW:$STY@$origin)";
      source .prompt
    fi
  else
    echo "You are on screen $WINDOW! ";
  fi
  local rm=`rm -f .screenOn-*`
}
autoScreenMe
alias as='autoScreenMe'
alias ms='myScreen'
alias mss='myScreens'

ssh() { 
  if [[ "$STY" != "" ]]; then 
    local timestamp=$(( `date "+%s"` + 120 ))
    local dateFormat="+%y%m%d%H%M.%S"
    local future=`date --date "1970-01-01 $timestamp sec" "$dateFormat"`
    # Get local information
    getMyIp;
    local sharedFile="$sharedFolder/.screenOn-$IP"
    # Get remote information
    getDomainName "$@";
    screen -S $STY -X title $DOMAIN_NAME &> /dev/null # Set title to remote
    echo "$STY:$WINDOW" > $sharedFile # Dial in the Stargate
    touch -t $future $sharedFile # Make shared file 120 seconds in the future
    command ssh "$@"; # Jump
    screen -S $STY -X title `hostname` &> /dev/null # Set title back to local
  else
    command ssh "$@"; # There is no spoon.
  fi
}
screen() { # If you must nest screens...
  export STYb=$STY
  export WINDOW=$WINDOW
  STY=""
  WINDOW=""
  command screen "$@"
  STY=$STYb
  WINDOW=$WINDOWb
}
cd2() {
  echo $@ > ~/.cd
  cd "$@"
}
